import csv
import os
from random import randint, random
from matplotlib import pyplot as plt
import os


from matplotlib import pyplot as plt
from tqdm import tqdm,trange

def draw(my_dict):
    titlename = "dataset information"
    plt.figure(figsize=(10, 5))
    plt.title(titlename)
    plt.grid(axis='y')
    lang = []
    cnt = []
    for k in sorted(my_dict):
        lang.append(k)
        cnt.append(len(my_dict[k]))
    plt.bar(lang, cnt, align="center", color="steelblue", alpha=0.6)
    plt.xlabel("lang")
    plt.ylabel("count")
    plt.legend()
    # if idx % 20 == 0:
    plt.savefig("data.png")

def get_data(in_name):
    fin = open(in_name)
    data = []
    for i, line in tqdm(enumerate(fin.readlines()[1:]), desc="get data"):
        l_split = line.strip().split('\t')
        p, h, l = l_split
        p = p.replace(',', ' ')
        p = p.replace('"', '')
        h = h.replace(',', ' ')
        h = h.replace('"', '')
        data.append((p, h, l))
    return data


def merge(all_data):
    t2cid = {}
    out_data = []
    all_lg = list(all_data.keys())
    for i, d in enumerate(all_data['en']):
        if d[2] == 'contradictory':  # hard negative
            t2cid[d[0]] = i
    error = 0
    for i, d in tqdm(enumerate(all_data['en']), desc="merge data"):
        if d[0] in t2cid:
            cid = t2cid[d[0]]
        else:
            cid = random.choice(list(t2cid.values()))
            error += 1
        if d[2] == 'entailment':
            for lg in all_lg:
                slg = random.choice(all_lg)
                slg2 = random.choice(all_lg)
                out_data.append(all_data[lg][i][0] + ',' + all_data[slg][i][1] + ',' + all_data[slg2][cid][1])
    print('error:', error)
    return out_data


data_dir = 'data/XNLI-MT-1.0/multinli/'
all_data = {}
for f in os.listdir('data/XNLI-MT-1.0/multinli/'):
    lg = f.split('.')[2]
    all_data[lg] = get_data(data_dir + f)
draw(all_data)
out_data = merge(all_data)
print(out_data[10])

# prepared dataset according to CoSDA-ML

# from prepared import read_dict,build





fin = open('data/nli_for_simcse.csv')
for line in fin.readlines()[1:]:
    out_data.append(line.strip())
print(out_data[-1])

random.shuffle(out_data)
fout = open('data/nli_merged_all.csv', 'w')
fout.write('sent0,sent1,hard_neg\n')
fout.write('\n'.join(out_data))
